package com.hlz.flink.chapter05;

import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

/**
 * rebalance 是所有分区数据的“重新平衡”
 * <p>
 * 用 rescale 的方式进行“局部重缩放”，就可以让数据只在当前 TaskManager 的多个 slot 之间重新分配，从而避免了网络
 * 传输带来的损耗。
 *
 * @author Hongliang Zhu
 * @create 2023-05-06 23:27
 */
public class RescalePartitionTest {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment executionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment();

        executionEnvironment.setParallelism(1);

        // 使用并行数据源的富函数，这样可以调用getRuntimeContext方法来执行运行时上下文的信息
        executionEnvironment.addSource(new RichParallelSourceFunction<Integer>() {
            @Override
            public void run(SourceContext<Integer> ctx) throws Exception {
                for (int i = 0; i < 8; i++) {
                    // 将奇数发送到索引为1的并行子任务
                    // 偶数发送到索引为0的并行子任务
                    if ((i + 1) % 2 == getRuntimeContext().getIndexOfThisSubtask()) {
                        ctx.collect(i + 1);
                    }

                }
            }

            @Override
            public void cancel() {

            }
        }).setParallelism(2)
                .rebalance()
                .print("rescale")
                .setParallelism(4);


        executionEnvironment.execute();

    }
}
